import socket
import json
import uuid
import time
import zmq
import threading
from typing import Dict
from common import *

class ThreadDeviceAgent:
    def __init__(self, name: str, info: str = "", pub_port: int = 6000):
        self.uuid = str(uuid.uuid4())
        self.name = name
        self.info = info
        self.pub_port = pub_port
        self.ip = get_local_ip()
        self.context = zmq.Context()
        self.running = True
        self.threads = []

    def udp_broadcast(self):
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        message = json.dumps({
            "uuid": self.uuid,
            "name": self.name,
            "info": self.info,
            "pub_port": self.pub_port
        }).encode()

        while self.running:
            udp_socket.sendto(message, ('255.255.255.255', BROADCAST_PORT))
            time.sleep(BROADCAST_INTERVAL)

    def zmq_heartbeat(self):
        pub_socket = self.context.socket(zmq.PUB)
        pub_socket.bind(f"tcp://*:{self.pub_port}")

        while self.running:
            msg = json.dumps({
                "uuid": self.uuid,
                "timestamp": time.time()
            })
            pub_socket.send_string(msg)
            time.sleep(HEARTBEAT_INTERVAL)

    def start(self):
        print(f"[Agent] 启动设备 {self.name}（UUID: {self.uuid}）")
        retry_count = 0

        while retry_count < MAX_RETRY_COUNT and self.running:
            try:
                # 创建并启动线程
                broadcast_thread = threading.Thread(target=self.udp_broadcast)
                heartbeat_thread = threading.Thread(target=self.zmq_heartbeat)
                
                self.threads = [broadcast_thread, heartbeat_thread]
                
                for thread in self.threads:
                    thread.daemon = True
                    thread.start()
                
                # 等待线程结束
                for thread in self.threads:
                    thread.join()
                    
            except Exception as e:
                retry_count += 1
                print(f"[错误] 设备运行异常: {e}")
                if retry_count < MAX_RETRY_COUNT:
                    print(f"[重试] {RECONNECT_DELAY}秒后进行第{retry_count}次重试...")
                    time.sleep(RECONNECT_DELAY)
                else:
                    print(f"[失败] 已达到最大重试次数({MAX_RETRY_COUNT})，设备停止运行")
                    break

    def stop(self):
        self.running = False
        # 等待所有线程结束
        for thread in self.threads:
            thread.join(timeout=5)

if __name__ == "__main__":
    agent = ThreadDeviceAgent(name="AlphaBot", info="Jetson Nano")
    try:
        agent.start()
    except KeyboardInterrupt:
        agent.stop()